import os
import json
import re
import math
from PIL import Image

root_dir = r"D:\desktop\road_detection\workspace\result\bev"
scale_length_m = 20.0
scale_px_length = 100 # 121.0
snipaste_px_per_m = scale_px_length / scale_length_m
bev_px_per_m = 20.0

def deg2m(lat):
    return 111320.0, 111320.0 * math.cos(math.radians(lat))

for entry in os.scandir(root_dir):
    if not entry.is_dir():
        continue
    folder = entry.path
    name = os.path.basename(folder)
    # if name not in ['data331_b9']:
    #     continue
    pngs = [f for f in os.listdir(folder) if f.startswith("Snipaste") and f.lower().endswith(".png")]
    if len(pngs) != 1:
        continue

    txts = [f for f in os.listdir(folder) if f.lower().endswith(".txt")]
    if len(txts) != 1:
        continue

    jsons = [f for f in os.listdir(folder) if f.lower().endswith("meta.json")]
    if not jsons:
        jsons = [f for f in os.listdir(folder) if f.lower().endswith(".json")]
    if not jsons:
        continue

    cfg = json.load(open(os.path.join(folder, jsons[0]), 'r', encoding='utf-8'))
    origin_lat, origin_lon = cfg['lat0_lon0']
    ll_mx, ll_my = cfg['left_bottom']
    ur_mx, ur_my = cfg['right_top']

    width_m = ur_mx - ll_mx
    height_m = ur_my - ll_my
    out_w = int(round(width_m * bev_px_per_m))
    out_h = int(round(height_m * bev_px_per_m))
    canvas = Image.new("RGB", (out_w, out_h))

    m_lat, m_lon = deg2m(origin_lat)

    png = pngs[0]
    txt_path = os.path.join(folder, txts[0])
    with open(txt_path, 'r', encoding='utf-8') as f:
        txt = f.read()
    m = re.search(r'@([0-9\.\-]+),([0-9\.\-]+)', txt)
    if not m:
        continue
    center_lat, center_lon = map(float, m.groups())

    dx_m = (center_lon - origin_lon) * m_lon
    dy_m = (center_lat - origin_lat) * m_lat

    img = Image.open(os.path.join(folder, png))
    w, h = img.size

    scale = bev_px_per_m / snipaste_px_per_m
    new_w = int(round(w * scale))
    new_h = int(round(h * scale))
    img = img.resize((new_w, new_h), Image.LANCZOS)

    px_center = int(round((dx_m - ll_mx) * bev_px_per_m))
    py_center = int(round((ur_my - dy_m) * bev_px_per_m))

    x0 = px_center - new_w // 2
    y0 = py_center - new_h // 2

    canvas.paste(img, (x0, y0))

    out_file = os.path.join(folder, f"{name}_mosaic.png")
    canvas.save(out_file)
